import os
import numpy as np
import cv2

def load_flow(flow_path):
    return np.load(flow_path)

def save_flow(flow, output_path):
    np.save(output_path, flow)
    print(f"Saved {output_path}")


def composite_flows(flow_a, flow_b):
    h, w, _ = flow_a.shape

    # 创建映射矩阵
    flow_a_map = np.zeros((h, w, 2), dtype=np.float32)
    flow_b_map = np.zeros((h, w, 2), dtype=np.float32)

    # 用光流场 A 和 B 的偏移创建映射矩阵
    flow_a_map[..., 0] = flow_a[..., 0] + np.arange(w)
    flow_a_map[..., 1] = flow_a[..., 1] + np.arange(h)[:, np.newaxis]

    flow_b_map[..., 0] = flow_b[..., 0] + np.arange(w)
    flow_b_map[..., 1] = flow_b[..., 1] + np.arange(h)[:, np.newaxis]

    # 将光流场 A 应用于光流场 B 的偏移
    flow_a_remap_x = cv2.remap(flow_a[..., 0], flow_b_map[..., 0], flow_b_map[..., 1], interpolation=cv2.INTER_CUBIC)
    flow_a_remap_y = cv2.remap(flow_a[..., 1], flow_b_map[..., 0], flow_b_map[..., 1], interpolation=cv2.INTER_CUBIC)

    # 计算复合光流场
    flow_c = np.zeros_like(flow_a)
    flow_c[..., 0] = flow_b[..., 0] + flow_a_remap_x
    flow_c[..., 1] = flow_b[..., 1] + flow_a_remap_y

    return flow_c

def process_and_save_composite_flows(flow_a_dir, flow_b_dir, output_dir):
    os.makedirs(output_dir, exist_ok=True)

    flow_a_files = sorted([f for f in os.listdir(flow_a_dir) if f.endswith('.npy')])
    flow_b_files = sorted([f for f in os.listdir(flow_b_dir) if f.endswith('.npy')])

    common_files = set(flow_a_files).intersection(set(flow_b_files))

    for filename in common_files:
        flow_a_path = os.path.join(flow_a_dir, filename)
        flow_b_path = os.path.join(flow_b_dir, filename)
        output_path = os.path.join(output_dir, filename)

        flow_a = load_flow(flow_a_path)
        flow_b = load_flow(flow_b_path)
        flow_c = composite_flows(flow_a, flow_b)

        save_flow(flow_c, output_path)

# 示例用法
flow_a_dir = '/root/data/deep_align/coarse_32nm_big_flow2'  # 替换为实际的A光流场文件夹路径
flow_b_dir = '/root/data/deep_align/32nm_flow_big_fine_align(3)'  # 替换为实际的B光流场文件夹路径
output_dir = '/root/data/deep_align/32nm_flow_big(3)'  # 替换为实际输出文件夹路径

process_and_save_composite_flows(flow_a_dir, flow_b_dir, output_dir)
